# -*- coding: utf-8 -*-
"""
数据源服务 - 统一管理mysql和starrocks数据源连接
"""
import json
import time
import logging
import pymysql
import pymysql.cursors
from typing import Dict, List, Optional, Any
from contextlib import contextmanager

from app.utils.mysql_db import get_db_connection

logger = logging.getLogger(__name__)


class DataSourceService:
    """数据源服务 - 统一管理mysql和starrocks连接"""
    
    # 支持的数据源类型
    SUPPORTED_TYPES = ['mysql', 'starrocks']
    
    @staticmethod
    def get_config(config_id: int) -> Dict[str, Any]:
        """
        获取数据源配置
        :param config_id: service_config表的ID
        :return: 配置字典
        """
        conn = get_db_connection()
        cursor = conn.cursor(pymysql.cursors.DictCursor)
        
        try:
            sql = """
                SELECT id, service_type, config_name, host, port, username, password, 
                       `database`, extra_config, description, is_active, is_deleted
                FROM service_config 
                WHERE id = %s AND service_type IN ('mysql', 'starrocks')
                AND is_deleted = 0 AND is_active = 1
            """
            cursor.execute(sql, (config_id,))
            config = cursor.fetchone()
            
            if not config:
                raise ValueError(f"数据源配置不存在或未激活: config_id={config_id}")
            
            # 验证service_type
            if config['service_type'] not in DataSourceService.SUPPORTED_TYPES:
                raise ValueError(
                    f"不支持的数据源类型: {config['service_type']}, "
                    f"仅支持: {', '.join(DataSourceService.SUPPORTED_TYPES)}"
                )
            
            # 解析extra_config
            if config.get('extra_config'):
                try:
                    extra = json.loads(config['extra_config']) if isinstance(config['extra_config'], str) else config['extra_config']
                    config.update(extra)
                except (json.JSONDecodeError, TypeError):
                    logger.warning(f"解析extra_config失败: config_id={config_id}")
            
            return config
        finally:
            cursor.close()
            conn.close()
    
    @staticmethod
    def get_configs(service_type: Optional[str] = None, 
                   config_ids: Optional[List[int]] = None) -> List[Dict[str, Any]]:
        """
        获取数据源配置列表
        :param service_type: 服务类型（mysql或starrocks），None表示全部
        :param config_ids: 配置ID列表，None表示全部
        :return: 配置列表
        """
        conn = get_db_connection()
        cursor = conn.cursor(pymysql.cursors.DictCursor)
        
        try:
            conditions = ["service_type IN ('mysql', 'starrocks')", "is_deleted = 0", "is_active = 1"]
            params = []
            
            if service_type:
                if service_type not in DataSourceService.SUPPORTED_TYPES:
                    raise ValueError(f"不支持的数据源类型: {service_type}")
                conditions.append("service_type = %s")
                params.append(service_type)
            
            if config_ids:
                placeholders = ','.join(['%s'] * len(config_ids))
                conditions.append(f"id IN ({placeholders})")
                params.extend(config_ids)
            
            sql = f"""
                SELECT id, service_type, config_name, host, port, username, password, 
                       `database`, extra_config, description
                FROM service_config 
                WHERE {' AND '.join(conditions)}
                ORDER BY service_type, sort_order, id
            """
            cursor.execute(sql, params)
            configs = cursor.fetchall()
            
            # 解析extra_config
            for config in configs:
                if config.get('extra_config'):
                    try:
                        extra = json.loads(config['extra_config']) if isinstance(config['extra_config'], str) else config['extra_config']
                        config.update(extra)
                    except (json.JSONDecodeError, TypeError):
                        pass
            
            return configs
        finally:
            cursor.close()
            conn.close()
    
    @staticmethod
    def create_connection(config_id: int, autocommit: bool = False) -> pymysql.Connection:
        """
        创建数据库连接（mysql和starrocks都使用pymysql）
        :param config_id: service_config表的ID
        :param autocommit: 是否自动提交
        :return: 数据库连接
        """
        config = DataSourceService.get_config(config_id)
        
        connection_params = {
            'host': config['host'],
            'port': config['port'],
            'user': config['username'],
            'password': config['password'],
            'database': config.get('database', ''),
            'charset': config.get('charset', 'utf8mb4'),
            'cursorclass': pymysql.cursors.DictCursor,
            'autocommit': autocommit,
            'connect_timeout': config.get('connect_timeout', 10),
            'read_timeout': config.get('read_timeout', 30),
            'write_timeout': config.get('write_timeout', 30)
        }
        
        # 添加重试逻辑
        max_retries = 3
        last_error = None
        
        for retry in range(max_retries):
            try:
                if retry > 0:
                    logger.info(
                        f"🔄 重试{config['service_type']}连接 "
                        f"(尝试{retry+1}/{max_retries}) [{config['config_name']}]"
                    )
                
                conn = pymysql.connect(**connection_params)
                logger.debug(
                    f"{config['service_type']}连接创建成功: "
                    f"{config['config_name']} ({config['host']}:{config['port']})"
                )
                return conn
                
            except Exception as e:
                last_error = e
                logger.error(
                    f"创建{config['service_type']}连接失败 "
                    f"(尝试{retry+1}/{max_retries}) [{config['config_name']}]: {e}"
                )
                
                if retry < max_retries - 1:
                    # 递增等待：第1次失败等30秒，第2次失败等60秒
                    wait_time = 30 if retry == 0 else 60
                    logger.info(f"⏰ 等待{wait_time}秒后重试...")
                    time.sleep(wait_time)
                else:
                    logger.error(f"❌ {config['service_type']}连接重试{max_retries}次均失败")
        
        # 最终失败，抛出最后的异常
        raise last_error
    
    @staticmethod
    @contextmanager
    def get_connection(config_id: int, autocommit: bool = False):
        """
        获取数据库连接的上下文管理器
        :param config_id: service_config表的ID
        :param autocommit: 是否自动提交
        :yield: 数据库连接
        """
        conn = None
        try:
            conn = DataSourceService.create_connection(config_id, autocommit)
            yield conn
            if not autocommit:
                conn.commit()
        except Exception as e:
            if conn:
                conn.rollback()
            logger.error(f"数据库操作失败: {e}")
            raise
        finally:
            if conn:
                conn.close()
                logger.debug("数据库连接已关闭")
    
    @staticmethod
    @contextmanager
    def get_cursor(config_id: int, autocommit: bool = False):
        """
        获取数据库游标的上下文管理器
        :param config_id: service_config表的ID
        :param autocommit: 是否自动提交
        :yield: 数据库游标
        """
        with DataSourceService.get_connection(config_id, autocommit) as conn:
            cursor = conn.cursor()
            try:
                yield cursor
            finally:
                cursor.close()
    
    @staticmethod
    def test_connection(config_id: int) -> Dict[str, Any]:
        """
        测试数据源连接
        :param config_id: service_config表的ID
        :return: 测试结果
        """
        try:
            config = DataSourceService.get_config(config_id)
            
            # 尝试创建连接
            conn = DataSourceService.create_connection(config_id)
            conn.close()
            
            return {
                'success': True,
                'message': f"连接成功: {config['config_name']} ({config['host']}:{config['port']})",
                'config': {
                    'id': config['id'],
                    'service_type': config['service_type'],
                    'config_name': config['config_name'],
                    'host': config['host'],
                    'port': config['port']
                }
            }
        except Exception as e:
            return {
                'success': False,
                'message': f"连接失败: {str(e)}",
                'error': str(e)
            }

